"""
物理计划优化器
将逻辑执行计划转换为物理执行计划（Engine可执行的QueryPlan）
"""

from typing import List, Dict, Any, Optional
import sys
import os

# 导入逻辑计划相关
from logical_plan_generator import LogicalNode, LogicalOperatorType

# 导入Engine相关
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from ExecutionEngine.query_plan import QueryPlan, PlanNode, PlanBuilder
from ExecutionEngine.types import OperatorType, JoinType


class PhysicalPlanOptimizer:
    """物理计划优化器"""
    
    def __init__(self):
        self.plan_builder = PlanBuilder()
        self.node_counter = 0
    
    def optimize_to_physical_plan(self, logical_root: LogicalNode) -> QueryPlan:
        """
        将逻辑计划优化并转换为物理执行计划
        
        Args:
            logical_root: 逻辑计划的根节点
            
        Returns:
            QueryPlan: 可执行的查询计划
        """
        # 转换逻辑节点为物理节点
        physical_root = self._convert_logical_to_physical(logical_root)
        
        # 构建查询计划
        query_plan = self.plan_builder.build(physical_root)
        
        # 应用物理优化
        self._apply_physical_optimizations(query_plan)
        
        return query_plan
    
    def _convert_logical_to_physical(self, logical_node: LogicalNode) -> PlanNode:
        """将逻辑节点转换为物理节点"""
        
        # 递归转换子节点
        physical_children = []
        for child in logical_node.children:
            physical_child = self._convert_logical_to_physical(child)
            physical_children.append(physical_child)
        
        # 根据逻辑操作符类型创建对应的物理节点
        if logical_node.operator_type == LogicalOperatorType.SCAN:
            return self._create_scan_node(logical_node, physical_children)
        
        elif logical_node.operator_type == LogicalOperatorType.PROJECT:
            return self._create_project_node(logical_node, physical_children)
        
        elif logical_node.operator_type == LogicalOperatorType.FILTER:
            return self._create_filter_node(logical_node, physical_children)
        
        elif logical_node.operator_type == LogicalOperatorType.JOIN:
            return self._create_join_node(logical_node, physical_children)
        
        elif logical_node.operator_type == LogicalOperatorType.AGGREGATE:
            return self._create_aggregate_node(logical_node, physical_children)
        
        elif logical_node.operator_type == LogicalOperatorType.SORT:
            return self._create_sort_node(logical_node, physical_children)
        
        elif logical_node.operator_type == LogicalOperatorType.LIMIT:
            return self._create_limit_node(logical_node, physical_children)
        
        elif logical_node.operator_type == LogicalOperatorType.CREATE_TABLE:
            return self._create_create_table_node(logical_node, physical_children)
        
        elif logical_node.operator_type == LogicalOperatorType.INSERT:
            return self._create_insert_node(logical_node, physical_children)
        
        elif logical_node.operator_type == LogicalOperatorType.DELETE:
            return self._create_delete_node(logical_node, physical_children)
        
        elif logical_node.operator_type == LogicalOperatorType.UPDATE:
            return self._create_update_node(logical_node, physical_children)
        
        elif logical_node.operator_type == LogicalOperatorType.UNION:
            return self._create_union_node(logical_node, physical_children)
        
        elif logical_node.operator_type == LogicalOperatorType.INTERSECT:
            return self._create_intersect_node(logical_node, physical_children)
        
        elif logical_node.operator_type == LogicalOperatorType.EXCEPT:
            return self._create_except_node(logical_node, physical_children)
        
        elif logical_node.operator_type == LogicalOperatorType.GRANT:
            return self._create_grant_node(logical_node, physical_children)
        
        elif logical_node.operator_type == LogicalOperatorType.CREATE_DATABASE:
            return self._create_create_database_node(logical_node, physical_children)
        
        elif logical_node.operator_type == LogicalOperatorType.DROP_DATABASE:
            return self._create_drop_database_node(logical_node, physical_children)
        
        elif logical_node.operator_type == LogicalOperatorType.USE_DATABASE:
            return self._create_use_database_node(logical_node, physical_children)
        
        elif logical_node.operator_type == LogicalOperatorType.LIST_DATABASES:
            return self._create_list_databases_node(logical_node, physical_children)
        
        elif logical_node.operator_type == LogicalOperatorType.CREATE_USER:
            return self._create_create_user_node(logical_node, physical_children)
        
        elif logical_node.operator_type == LogicalOperatorType.DROP_USER:
            return self._create_drop_user_node(logical_node, physical_children)
        
        else:
            raise ValueError(f"Unsupported logical operator: {logical_node.operator_type}")
    
    def _create_scan_node(self, logical_node: LogicalNode, 
                         physical_children: List[PlanNode]) -> PlanNode:
        """创建扫描物理节点"""
        table_name = logical_node.properties.get('table_name')
        alias = logical_node.properties.get('alias', table_name)
        
        # 根据表大小和索引情况选择扫描方式
        scan_type = self._choose_scan_type(table_name)
        
        physical_node = PlanNode(
            node_id=f"scan_{self.node_counter}",
            operator_type=scan_type,
            children=physical_children
        )
        
        physical_node.set_property('table_name', table_name)
        if alias != table_name:
            physical_node.set_property('alias', alias)
        
        self.node_counter += 1
        return physical_node
    
    def _create_project_node(self, logical_node: LogicalNode,
                           physical_children: List[PlanNode]) -> PlanNode:
        """创建投影物理节点"""
        columns = logical_node.properties.get('columns', [])
        
        physical_node = PlanNode(
            node_id=f"project_{self.node_counter}",
            operator_type=OperatorType.PROJECT,
            children=physical_children
        )
        
        physical_node.set_property('columns', columns)
        
        self.node_counter += 1
        return physical_node
    
    def _create_filter_node(self, logical_node: LogicalNode,
                          physical_children: List[PlanNode]) -> PlanNode:
        """创建过滤物理节点"""
        predicate = logical_node.properties.get('predicate')
        
        physical_node = PlanNode(
            node_id=f"filter_{self.node_counter}",
            operator_type=OperatorType.FILTER,
            children=physical_children
        )
        
        physical_node.set_property('predicate', predicate)
        
        self.node_counter += 1
        return physical_node
    
    def _create_join_node(self, logical_node: LogicalNode,
                         physical_children: List[PlanNode]) -> PlanNode:
        """创建连接物理节点"""
        join_condition = logical_node.properties.get('join_condition')
        logical_join_type = logical_node.properties.get('join_type', 'INNER')
        
        # 转换字符串JOIN类型到枚举
        join_type_mapping = {
            'INNER': JoinType.INNER,
            'LEFT': JoinType.LEFT,
            'RIGHT': JoinType.RIGHT,
            'FULL': JoinType.FULL,
            'CROSS': JoinType.CROSS
        }
        
        # 确保join_type是JoinType枚举
        if isinstance(logical_join_type, str):
            join_type = join_type_mapping.get(logical_join_type.upper(), JoinType.INNER)
        else:
            join_type = logical_join_type
        
        # 选择最优的连接算法
        join_algorithm = self._choose_join_algorithm(physical_children)
        
        physical_node = PlanNode(
            node_id=f"join_{self.node_counter}",
            operator_type=join_algorithm,
            children=physical_children
        )
        
        physical_node.set_property('join_condition', join_condition)
        physical_node.set_property('join_type', join_type)
        
        self.node_counter += 1
        return physical_node
    
    def _create_aggregate_node(self, logical_node: LogicalNode,
                             physical_children: List[PlanNode]) -> PlanNode:
        """创建聚合物理节点"""
        group_by = logical_node.properties.get('group_by', [])
        aggregates = logical_node.properties.get('aggregates', [])
        
        # 选择聚合算法
        agg_algorithm = self._choose_aggregate_algorithm(group_by, aggregates)
        
        physical_node = PlanNode(
            node_id=f"aggregate_{self.node_counter}",
            operator_type=agg_algorithm,
            children=physical_children
        )
        
        physical_node.set_property('group_by', group_by)
        physical_node.set_property('aggregates', aggregates)
        
        self.node_counter += 1
        return physical_node
    
    def _create_sort_node(self, logical_node: LogicalNode,
                         physical_children: List[PlanNode]) -> PlanNode:
        """创建排序物理节点"""
        order_by = logical_node.properties.get('order_by', [])
        
        physical_node = PlanNode(
            node_id=f"sort_{self.node_counter}",
            operator_type=OperatorType.SORT,
            children=physical_children
        )
        
        physical_node.set_property('order_by', order_by)
        
        self.node_counter += 1
        return physical_node
    
    def _create_limit_node(self, logical_node: LogicalNode,
                          physical_children: List[PlanNode]) -> PlanNode:
        """创建限制物理节点"""
        # 逻辑节点使用'count'属性，物理节点使用'limit'属性
        limit = logical_node.properties.get('count')  # 改为获取'count'属性
        offset = logical_node.properties.get('offset', 0)
        
        physical_node = PlanNode(
            node_id=f"limit_{self.node_counter}",
            operator_type=OperatorType.LIMIT,
            children=physical_children
        )
        
        physical_node.set_property('limit', limit)
        physical_node.set_property('offset', offset)
        
        self.node_counter += 1
        return physical_node
    
    def _choose_scan_type(self, table_name: str) -> OperatorType:
        """选择最优的扫描方式"""
        # 简化的选择逻辑，实际应该基于统计信息
        # 这里默认使用表扫描
        return OperatorType.TABLE_SCAN
    
    def _choose_join_algorithm(self, children: List[PlanNode]) -> OperatorType:
        """选择最优的连接算法"""
        # 简化的选择逻辑
        # 实际应该基于表大小、索引、内存等因素
        if len(children) == 2:
            # 默认使用哈希连接
            return OperatorType.HASH_JOIN
        else:
            raise ValueError("Join requires exactly 2 children")
    
    def _choose_aggregate_algorithm(self, group_by: List[str], 
                                  aggregates: List[Dict]) -> OperatorType:
        """选择最优的聚合算法"""
        # 简化的选择逻辑
        if group_by:
            # 有分组，使用哈希聚合
            return OperatorType.HASH_AGGREGATE
        else:
            # 无分组，使用哈希聚合
            return OperatorType.HASH_AGGREGATE
    
    def _apply_physical_optimizations(self, query_plan: QueryPlan):
        """应用物理优化"""
        # 应用各种物理优化规则
        optimization_rules = [
            'push_down_predicates',
            'push_down_projections',
            'join_reordering',
            'index_selection'
        ]
        
        query_plan.optimize(optimization_rules)
    
    def _create_create_table_node(self, logical_node: LogicalNode,
                                 physical_children: List[PlanNode]) -> PlanNode:
        """创建CREATE TABLE物理节点"""
        table_name = logical_node.properties.get('table_name')
        columns = logical_node.properties.get('columns', [])
        
        physical_node = PlanNode(
            node_id=f"create_table_{self.node_counter}",
            operator_type=OperatorType.CREATE_TABLE,
            children=physical_children
        )
        
        physical_node.set_property('table_name', table_name)
        physical_node.set_property('columns', columns)
        
        self.node_counter += 1
        return physical_node
    
    def _create_insert_node(self, logical_node: LogicalNode,
                           physical_children: List[PlanNode]) -> PlanNode:
        """创建INSERT物理节点"""
        table_name = logical_node.properties.get('table_name')
        values = logical_node.properties.get('values', [])
        
        physical_node = PlanNode(
            node_id=f"insert_{self.node_counter}",
            operator_type=OperatorType.INSERT,
            children=physical_children
        )
        
        physical_node.set_property('table_name', table_name)
        physical_node.set_property('values', values)
        
        self.node_counter += 1
        return physical_node
    
    def _create_update_node(self, logical_node: LogicalNode,
                           physical_children: List[PlanNode]) -> PlanNode:
        """创建UPDATE物理节点"""
        table_name = logical_node.properties.get('table_name')
        set_clauses = logical_node.properties.get('set_clauses', [])
        where_condition = logical_node.properties.get('where_condition')
        
        physical_node = PlanNode(
            node_id=f"update_{self.node_counter}",
            operator_type=OperatorType.UPDATE,
            children=physical_children
        )
        
        physical_node.set_property('table_name', table_name)
        physical_node.set_property('set_clauses', set_clauses)
        if where_condition:
            physical_node.set_property('where_condition', where_condition)
        
        self.node_counter += 1
        return physical_node
    
    def _create_delete_node(self, logical_node: LogicalNode,
                           physical_children: List[PlanNode]) -> PlanNode:
        """创建DELETE物理节点"""
        table_name = logical_node.properties.get('table_name')
        where_condition = logical_node.properties.get('where_condition')
        
        physical_node = PlanNode(
            node_id=f"delete_{self.node_counter}",
            operator_type=OperatorType.DELETE,
            children=physical_children
        )
        
        physical_node.set_property('table_name', table_name)
        physical_node.set_property('where_condition', where_condition)
        
        # 设置估算成本和行数
        physical_node.cost = 1.5
        physical_node.estimated_rows = 1
        
        self.node_counter += 1
        return physical_node
    
    def _create_union_node(self, logical_node: LogicalNode,
                          physical_children: List[PlanNode]) -> PlanNode:
        """创建UNION物理节点"""
        distinct = logical_node.properties.get('distinct', True)
        
        physical_node = PlanNode(
            node_id=f"union_{self.node_counter}",
            operator_type=OperatorType.UNION,
            children=physical_children
        )
        
        physical_node.set_property('distinct', distinct)
        
        self.node_counter += 1
        return physical_node
    
    def _create_intersect_node(self, logical_node: LogicalNode,
                              physical_children: List[PlanNode]) -> PlanNode:
        """创建INTERSECT物理节点"""
        distinct = logical_node.properties.get('distinct', True)
        
        physical_node = PlanNode(
            node_id=f"intersect_{self.node_counter}",
            operator_type=OperatorType.INTERSECT,
            children=physical_children
        )
        
        physical_node.set_property('distinct', distinct)
        
        self.node_counter += 1
        return physical_node
    
    def _create_except_node(self, logical_node: LogicalNode,
                           physical_children: List[PlanNode]) -> PlanNode:
        """创建EXCEPT物理节点"""
        distinct = logical_node.properties.get('distinct', True)
        
        physical_node = PlanNode(
            node_id=f"except_{self.node_counter}",
            operator_type=OperatorType.EXCEPT,
            children=physical_children
        )
        
        physical_node.set_property('distinct', distinct)
        
        self.node_counter += 1
        return physical_node
    
    def _create_grant_node(self, logical_node: LogicalNode,
                          physical_children: List[PlanNode]) -> PlanNode:
        """创建GRANT物理节点"""
        privileges = logical_node.properties.get('privileges')
        object_type = logical_node.properties.get('object_type')
        object_name = logical_node.properties.get('object_name')
        grantees = logical_node.properties.get('grantees')
        grant_option = logical_node.properties.get('grant_option')
        
        # 创建GRANT操作的物理节点
        physical_node = PlanNode(
            node_id=f"grant_{self.node_counter}",
            operator_type=OperatorType.GRANT,
            children=physical_children
        )
        
        physical_node.set_property('operation', 'GRANT')
        physical_node.set_property('privileges', privileges)
        physical_node.set_property('object_type', object_type)
        physical_node.set_property('object_name', object_name)
        physical_node.set_property('grantees', grantees)
        if grant_option:
            physical_node.set_property('grant_option', grant_option)
        
        self.node_counter += 1
        return physical_node
    
    def _create_create_database_node(self, logical_node: LogicalNode, 
                                    physical_children: List[PlanNode]) -> PlanNode:
        """创建CREATE DATABASE物理节点"""
        database_name = logical_node.properties.get('database_name')
        charset = logical_node.properties.get('charset', 'utf8')
        
        physical_node = PlanNode(
            node_id=f"create_database_{self.node_counter}",
            operator_type=OperatorType.CREATE_DATABASE,
            children=physical_children
        )
        
        physical_node.set_property('operation', 'CREATE_DATABASE')
        physical_node.set_property('database_name', database_name)
        physical_node.set_property('charset', charset)
        
        self.node_counter += 1
        return physical_node
    
    def _create_drop_database_node(self, logical_node: LogicalNode, 
                                  physical_children: List[PlanNode]) -> PlanNode:
        """创建DROP DATABASE物理节点"""
        database_name = logical_node.properties.get('database_name')
        force = logical_node.properties.get('force', False)
        
        physical_node = PlanNode(
            node_id=f"drop_database_{self.node_counter}",
            operator_type=OperatorType.DROP_DATABASE,
            children=physical_children
        )
        
        physical_node.set_property('operation', 'DROP_DATABASE')
        physical_node.set_property('database_name', database_name)
        physical_node.set_property('force', force)
        
        self.node_counter += 1
        return physical_node
    
    def _create_use_database_node(self, logical_node: LogicalNode, 
                                 physical_children: List[PlanNode]) -> PlanNode:
        """创建USE DATABASE物理节点"""
        database_name = logical_node.properties.get('database_name')
        
        physical_node = PlanNode(
            node_id=f"use_database_{self.node_counter}",
            operator_type=OperatorType.USE_DATABASE,
            children=physical_children
        )
        
        physical_node.set_property('operation', 'USE_DATABASE')
        physical_node.set_property('database_name', database_name)
        
        self.node_counter += 1
        return physical_node
    
    def _create_list_databases_node(self, logical_node: LogicalNode, 
                                   physical_children: List[PlanNode]) -> PlanNode:
        """创建LIST DATABASES物理节点"""
        physical_node = PlanNode(
            node_id=f"list_databases_{self.node_counter}",
            operator_type=OperatorType.LIST_DATABASES,
            children=physical_children
        )
        
        physical_node.set_property('operation', 'LIST_DATABASES')
        
        self.node_counter += 1
        return physical_node
    
    def _create_create_user_node(self, logical_node: LogicalNode,
                               physical_children: List[PlanNode]) -> PlanNode:
        """创建CREATE USER物理节点"""
        username = logical_node.properties.get('username')
        password = logical_node.properties.get('password')
        
        # 创建CREATE USER操作的物理节点
        physical_node = PlanNode(
            node_id=f"create_user_{self.node_counter}",
            operator_type=OperatorType.CREATE_USER,
            children=physical_children
        )
        
        physical_node.set_property('operation', 'CREATE_USER')
        physical_node.set_property('username', username)
        physical_node.set_property('password', password)
        
        self.node_counter += 1
        return physical_node
    
    def _create_drop_user_node(self, logical_node: LogicalNode,
                             physical_children: List[PlanNode]) -> PlanNode:
        """创建DROP USER物理节点"""
        username = logical_node.properties.get('username')
        
        # 创建DROP USER操作的物理节点
        physical_node = PlanNode(
            node_id=f"drop_user_{self.node_counter}",
            operator_type=OperatorType.DROP_USER,
            children=physical_children
        )
        
        physical_node.set_property('operation', 'DROP_USER')
        physical_node.set_property('username', username)
        
        self.node_counter += 1
        return physical_node
    
    def print_physical_plan(self, query_plan: QueryPlan):
        """打印物理执行计划"""
        print("物理执行计划:")
        print(query_plan.explain('text'))


# 测试函数
def test_physical_plan_optimization():
    """测试物理计划优化"""
    print("=" * 60)
    print("物理计划优化器测试")
    print("=" * 60)
    
    from logical_plan_generator import LogicalPlanGenerator, LogicalNode, LogicalOperatorType
    
    # 创建一个简单的逻辑计划
    scan_node = LogicalNode(
        operator_type=LogicalOperatorType.SCAN,
        properties={'table_name': 'users', 'alias': 'u'}
    )
    
    filter_node = LogicalNode(
        operator_type=LogicalOperatorType.FILTER,
        properties={'predicate': 'age > 25'},
        children=[scan_node]
    )
    
    project_node = LogicalNode(
        operator_type=LogicalOperatorType.PROJECT,
        properties={'columns': ['name', 'age']},
        children=[filter_node]
    )
    
    # 转换为物理计划
    optimizer = PhysicalPlanOptimizer()
    physical_plan = optimizer.optimize_to_physical_plan(project_node)
    
    # 打印结果
    optimizer.print_physical_plan(physical_plan)


if __name__ == "__main__":
    test_physical_plan_optimization()

